#include <mpi.h>
#include <stdio.h>
#include <unistd.h>
#include <algorithm>
#include <random>
#include <chrono>

#define PRINTR(format,...) \
if (rank==0) printf(format,##__VA_ARGS__)

#define PRINTI(format,...) \
printf("Proc#%d:\t" format,rank,##__VA_ARGS__)

#define PRINTRD(format,...) \
MPI_Barrier(MPI_COMM_WORLD); PRINTR(format,##__VA_ARGS__)

#define PRINTID(format,...) \
MPI_Barrier(MPI_COMM_WORLD); PRINTI(format,##__VA_ARGS__)

int main(int argc, char** argv){
    int rank, nproc;
    int TAG = 100;
    MPI_Request send_request, recv_request;
    MPI_Status send_status, recv_status;
    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &nproc);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    if (argc<2){
        PRINTR("usage: mpirun -np NUM_OF_PROCS %s NUM_OF_SERVERS\n", argv[0]);
        MPI_Finalize();
        exit(0);
    }
    int p;
    p = std::atoi(argv[1]);
    // PRINTI("p=%d\n",p);
    int* ranks_incl = new int[nproc];
    int index = 0;
    for (int i = rank%p; i<nproc; i+=p){
        ranks_incl[index++] = i;
    }
    MPI_Group world_group, worker_group, server_group;
    MPI_Comm worker_comm, server_comm;
    int worker_rank, server_rank;
    MPI_Comm_group(MPI_COMM_WORLD, &world_group);
    MPI_Group_incl(world_group, index, ranks_incl, &worker_group);
    MPI_Comm_create(MPI_COMM_WORLD, worker_group, &worker_comm);
    MPI_Comm_rank(worker_comm, &worker_rank);
    
    for (int i = 0; i<p; i++){
        ranks_incl[i] = i;
    }
    MPI_Group_incl(world_group, p, ranks_incl, &server_group);
    MPI_Comm_create(MPI_COMM_WORLD, server_group, &server_comm);
    if (rank<p){
        MPI_Comm_rank(server_comm, &server_rank);
    }

    std::default_random_engine e(std::chrono::high_resolution_clock::now().time_since_epoch().count()+rank);
    std::uniform_real_distribution<float> d(rank,rank+1);
    float data;
    float sum;

    int iter_cnt = 0;
    while (iter_cnt++ < 3){
        if (rank<p){
            data=0;
            MPI_Reduce(&data, &sum, 1, MPI_FLOAT, MPI_SUM, 0, worker_comm);
            // PRINTI("sum=%lf\n",sum);
            data=sum;
            MPI_Allreduce(&data, &sum, 1, MPI_FLOAT, MPI_SUM, server_comm);
            data = sum / (nproc-p);
            MPI_Bcast(&data, 1, MPI_FLOAT, 0, worker_comm);
        }
        else{
            data = d(e);
            // PRINTI("data=%lf\tflow:%d->%d\n",data,rank,rank%p);
            MPI_Reduce(&data, &sum, 1, MPI_FLOAT, MPI_SUM, 0, worker_comm);
            MPI_Bcast(&data,1,MPI_FLOAT,0,worker_comm);
            PRINTI("recv data=%lf\n",data);
        }
    }

    delete[] ranks_incl;
    MPI_Finalize();
}